package com.cu.langchain4j;

import com.cu.langchain4j.tools.Relation;
import dev.langchain4j.agent.tool.P;
import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.community.model.zhipu.ZhipuAiChatModel;
import dev.langchain4j.community.model.zhipu.ZhipuAiEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.Tokenizer;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.nio.file.FileSystems;
import java.nio.file.PathMatcher;
import java.time.Duration;
import java.util.Arrays;
import java.util.List;

public class TestMain {
    private static final String apiKey = System.getenv("apiKey");

    public static void main(String[] args) {
        // memory();

        //chat();
        //loadDocument();

        tools();
    }

    private static void tools() {
        ZhipuAiChatModel model = ZhipuAiChatModel.builder().apiKey("2696fd48800b443d8240ca1fd0ac8714.QPHi68awRf2qiMrb")
                .logRequests(true)
                .logResponses(true)
                .model("glm-4-flash")
                .temperature(0.6)
                .maxToken(1024)
                .maxRetries(1)
                .callTimeout(Duration.ofSeconds(60))
                .connectTimeout(Duration.ofSeconds(60))
                .writeTimeout(Duration.ofSeconds(60))
                .readTimeout(Duration.ofSeconds(60))
                .build();


        ChatMemory memory = MessageWindowChatMemory.withMaxMessages(10);
        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(model)
                .tools(new Relation())
                .chatMemory(memory).build();
        System.out.println(assistant.chat("小摘摘和金天是什么关系？"));
    }


    interface Assistant {

        String chat(String message);
    }

    /**
     * 测试记忆功能
     */
    public static void memory() {
        ZhipuAiChatModel model = ZhipuAiChatModel.builder().apiKey("2696fd48800b443d8240ca1fd0ac8714.QPHi68awRf2qiMrb")
                .logRequests(true)
                .logResponses(true)
                .model("glm-4-flash")
                .temperature(0.6)
                .maxToken(1024)
                .maxRetries(1)
                .callTimeout(Duration.ofSeconds(60))
                .connectTimeout(Duration.ofSeconds(60))
                .writeTimeout(Duration.ofSeconds(60))
                .readTimeout(Duration.ofSeconds(60))
                .build();


        ChatMemory memory = MessageWindowChatMemory.withMaxMessages(10);
        Assistant assistant = AiServices.builder(Assistant.class).chatLanguageModel(model).chatMemory(memory).build();
        System.out.println(assistant.chat("丁梦平是金天的老婆，金天经常称呼丁梦平为小摘摘。"));
        String answer = assistant.chat("小摘摘和金天是什么关系？");
        System.out.println(answer);
    }

    /**
     * 测试加载文档功能查询
     */
    public static void loadDocument() {
        // 加载文档
        List<Document> documents = FileSystemDocumentLoader.loadDocumentsRecursively("C:\\Users\\Administrator\\Desktop\\测试文档内容\\知识点\\框架");
        // 本地内存向量数据库  可以实例化其他的向量数据库embeddingStore
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        // 初始化文档分割和向量生成模型
        EmbeddingModel embeddingModel = ZhipuAiEmbeddingModel.builder().apiKey("2696fd48800b443d8240ca1fd0ac8714.QPHi68awRf2qiMrb").logRequests(true)
                .logResponses(true)
                .model("embedding-3") // 这里模型要使用embedding-3或者embedding-2 其他的模型不支持向量切割功能
                .maxRetries(1)
                .callTimeout(Duration.ofSeconds(60))
                .connectTimeout(Duration.ofSeconds(60))
                .writeTimeout(Duration.ofSeconds(60))
                .readTimeout(Duration.ofSeconds(60))
                .dimensions(512)
                .build();

        DocumentByParagraphSplitter documentByParagraphSplitter = new DocumentByParagraphSplitter(1000, 200);
        // 创建将文档分割存入向量数据库的实例
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(documentByParagraphSplitter)
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();

        // 将文档分割并存入向量数据库
        ingestor.ingest(documents);

        // 解析对应查询向量
        Embedding embedding = embeddingModel.embed("springboot如何集成打印日志").content();

        // 打印查询向量的长度
        System.out.println("Query vector length: " + embedding.vector().length);

        // 构建向量搜索请求
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(embedding)
                .maxResults(10)
                .minScore(0.85)
                .build();

        // 执行向量搜索并获取匹配结果
        List<EmbeddingMatch<TextSegment>> matches = embeddingStore.search(embeddingSearchRequest).matches();
        // 遍历并打印匹配结果
        for (EmbeddingMatch<TextSegment> match : matches) {
            System.out.println(match.embedded().text());
        }

    }

    /**
     * 测试聊天功能
     */
    public static void chat() {
        ZhipuAiChatModel model = ZhipuAiChatModel.builder().apiKey("2696fd48800b443d8240ca1fd0ac8714.QPHi68awRf2qiMrb")
                .logRequests(true)
                .logResponses(true)
                .model("glm-4-flash")
                .temperature(0.6)
                .maxToken(1024)
                .maxRetries(1)
                .callTimeout(Duration.ofSeconds(60))
                .connectTimeout(Duration.ofSeconds(60))
                .writeTimeout(Duration.ofSeconds(60))
                .readTimeout(Duration.ofSeconds(60))
                .build();

        UserMessage message1 = UserMessage.from("丁梦平是金天的老婆，金天经常称呼丁梦平为小摘摘。");
        AiMessage message2 = model.chat(message1).aiMessage();

        System.out.println(message2.toString());
        UserMessage secondUserMessage = UserMessage.from("金天和丁梦平是什么关系？");
        AiMessage secondAiMessage = model.chat(message1, message2, secondUserMessage).aiMessage();
        System.out.println(secondAiMessage.toString());
    }
}
